"""
    BalanceOperators

Advanced balance constraint operators for maintaining dynamical consistency in 
atmospheric background error covariance modeling. This module implements the 
sophisticated balance relationships used in operational GSI to ensure that 
background error perturbations respect fundamental atmospheric dynamics.

Balance constraints are essential for physical consistency in variational data 
assimilation, ensuring that analysis increments:
- Satisfy geostrophic balance: ∇×v = f/∇²ψ
- Maintain hydrostatic equilibrium: ∂Φ/∂p = -RT/p  
- Preserve thermal wind balance: ∂v_g/∂ln(p) = (R/f)×∇T
- Respect mass-wind relationships in atmospheric flows

Mathematical Framework:
The balance operator B_bal transforms unbalanced variables into balanced increments:
```
x_balanced = B_bal · x_unbalanced
```

Where the operator enforces multiple constraints simultaneously through:
- Linear balance for basic geostrophic relationships
- Non-linear balance for more complex atmospheric states
- Iterative balance for strongly nonlinear systems
- Weak constraint balance for approximate relationships

Key Applications:
- Background error covariance preconditioning
- Analysis increment post-processing  
- Ensemble perturbation balancing
- 4D-Var tangent linear model consistency
"""
module BalanceOperators

using LinearAlgebra
using SparseArrays
using StaticArrays

export AbstractBalanceOperator, GeostrophicBalance, HydrostaticBalance
export ThermalWindBalance, NonlinearBalance, WeakConstraintBalance
export apply_balance_operator, apply_balance_constraint, compute_balance_residual
export BalanceConfiguration, create_balance_operator, validate_balance

"""
    AbstractBalanceOperator{T<:AbstractFloat}

Abstract base type for balance constraint operators.
"""
abstract type AbstractBalanceOperator{T<:AbstractFloat} end

"""
    BalanceConfiguration{T<:AbstractFloat}

Configuration parameters for balance operators.

# Fields
- `grid_spacing::Tuple{T,T}`: Grid spacing (dx, dy) in meters
- `coriolis_parameter::Matrix{T}`: Coriolis parameter f = 2Ω sin(φ)
- `latitude::Matrix{T}`: Latitude field (radians)
- `longitude::Matrix{T}`: Longitude field (radians)  
- `reference_pressure::T`: Reference pressure level (Pa)
- `gas_constant::T`: Specific gas constant for dry air (J/kg/K)
- `gravity::T`: Gravitational acceleration (m/s²)
- `earth_radius::T`: Earth radius (m)
- `rotation_rate::T`: Earth rotation rate (rad/s)
"""
struct BalanceConfiguration{T<:AbstractFloat}
    grid_spacing::Tuple{T,T}
    coriolis_parameter::Matrix{T}
    latitude::Matrix{T}
    longitude::Matrix{T}
    reference_pressure::T
    gas_constant::T
    gravity::T
    earth_radius::T
    rotation_rate::T
    
    function BalanceConfiguration{T}(;
        grid_spacing::Tuple{T,T} = (T(50000), T(50000)),  # 50 km default
        latitude::Matrix{T} = reshape(T.(-90:2:90), 1, :) * π/180,  # Default lat grid
        longitude::Matrix{T} = reshape(T.(0:4:356), :, 1) * π/180,   # Default lon grid
        reference_pressure::T = T(100000),     # 1000 hPa
        gas_constant::T = T(287.05),          # J/kg/K
        gravity::T = T(9.80665),              # m/s²
        earth_radius::T = T(6.371e6),         # m
        rotation_rate::T = T(7.2921e-5)       # rad/s
    ) where T
        
        # Compute Coriolis parameter: f = 2Ω sin(φ)
        coriolis_parameter = 2 * rotation_rate .* sin.(latitude)
        
        new{T}(grid_spacing, coriolis_parameter, latitude, longitude,
               reference_pressure, gas_constant, gravity, earth_radius, rotation_rate)
    end
end

BalanceConfiguration(args...; kwargs...) = BalanceConfiguration{Float64}(args...; kwargs...)

"""
    GeostrophicBalance{T<:AbstractFloat} <: AbstractBalanceOperator{T}

Geostrophic balance constraint operator enforcing f×v = -∇Φ.

The geostrophic approximation assumes that the horizontal pressure gradient 
force balances the Coriolis force:
```
f×v_g = -∇Φ  ⟹  u_g = -(1/f)∂Φ/∂y, v_g = (1/f)∂Φ/∂x
```

# Fields
- `config::BalanceConfiguration{T}`: Balance configuration parameters
- `gradient_operators::Tuple{SparseMatrixCSC{T,Int}, SparseMatrixCSC{T,Int}}`: ∇x, ∇y operators
- `divergence_operator::SparseMatrixCSC{T,Int}`: Divergence operator ∇·
- `laplacian_operator::SparseMatrixCSC{T,Int}`: Laplacian operator ∇²
- `streamfunction_solver::LinearSolve.SciMLLinearSolveAlgorithm`: Efficient ψ solver
"""
struct GeostrophicBalance{T<:AbstractFloat} <: AbstractBalanceOperator{T}
    config::BalanceConfiguration{T}
    gradient_operators::Tuple{SparseMatrixCSC{T,Int}, SparseMatrixCSC{T,Int}}
    divergence_operator::SparseMatrixCSC{T,Int}
    laplacian_operator::SparseMatrixCSC{T,Int}
    
    function GeostrophicBalance{T}(config::BalanceConfiguration{T}) where T
        nx, ny = size(config.latitude)
        
        # Create finite difference operators
        grad_x, grad_y = create_gradient_operators(T, nx, ny, config.grid_spacing)
        div_op = create_divergence_operator(grad_x, grad_y)
        lapl_op = create_laplacian_operator(T, nx, ny, config.grid_spacing)
        
        new{T}(config, (grad_x, grad_y), div_op, lapl_op)
    end
end

GeostrophicBalance(config::BalanceConfiguration{T}) where T = GeostrophicBalance{T}(config)

"""
    HydrostaticBalance{T<:AbstractFloat} <: AbstractBalanceOperator{T}

Hydrostatic balance constraint operator enforcing ∂Φ/∂p = -RT/p.

The hydrostatic approximation assumes vertical pressure gradient force 
balances gravity:
```
∂p/∂z = -ρg  ⟹  ∂Φ/∂p = -RT/p
```

# Fields
- `config::BalanceConfiguration{T}`: Balance configuration
- `pressure_levels::Vector{T}`: Vertical pressure coordinate levels
- `vertical_derivative_operator::SparseMatrixCSC{T,Int}`: ∂/∂p operator
- `temperature_operator::SparseMatrixCSC{T,Int}`: Temperature interpolation
"""
struct HydrostaticBalance{T<:AbstractFloat} <: AbstractBalanceOperator{T}
    config::BalanceConfiguration{T}
    pressure_levels::Vector{T}
    vertical_derivative_operator::SparseMatrixCSC{T,Int}
    temperature_operator::SparseMatrixCSC{T,Int}
    
    function HydrostaticBalance{T}(config::BalanceConfiguration{T}, 
                                  pressure_levels::Vector{T}) where T
        n_levels = length(pressure_levels)
        
        # Create vertical derivative operator
        vert_deriv_op = create_vertical_derivative_operator(T, pressure_levels)
        
        # Temperature interpolation operator (simplified)
        temp_op = sparse(I, n_levels, n_levels)  # Identity for now
        
        new{T}(config, pressure_levels, vert_deriv_op, temp_op)
    end
end

"""
    ThermalWindBalance{T<:AbstractFloat} <: AbstractBalanceOperator{T}

Thermal wind balance constraint: ∂v_g/∂ln(p) = (R/f)k×∇T.

Relates vertical wind shear to horizontal temperature gradients:
```
∂u_g/∂ln(p) = -(R/f)∂T/∂y, ∂v_g/∂ln(p) = (R/f)∂T/∂x
```

# Fields  
- `config::BalanceConfiguration{T}`: Balance configuration
- `pressure_levels::Vector{T}`: Vertical coordinate levels
- `thermal_wind_operator::SparseMatrixCSC{T,Int}`: Combined thermal wind operator
"""
struct ThermalWindBalance{T<:AbstractFloat} <: AbstractBalanceOperator{T}
    config::BalanceConfiguration{T}
    pressure_levels::Vector{T}
    thermal_wind_operator::SparseMatrixCSC{T,Int}
    
    function ThermalWindBalance{T}(config::BalanceConfiguration{T},
                                  pressure_levels::Vector{T}) where T
        # Create thermal wind relationship operator
        thermal_op = create_thermal_wind_operator(T, config, pressure_levels)
        
        new{T}(config, pressure_levels, thermal_op)
    end
end

"""
    NonlinearBalance{T<:AbstractFloat} <: AbstractBalanceOperator{T}

Nonlinear balance constraint for strongly nonlinear atmospheric states.

Solves the nonlinear balance equation iteratively:
```
∇²ψ + f∇²χ + ∇f·∇χ = ∇×v + ...higher order terms
```

# Fields
- `linear_operator::GeostrophicBalance{T}`: Linear component
- `max_iterations::Int`: Maximum iterations for nonlinear solver
- `tolerance::T`: Convergence tolerance
- `relaxation_factor::T`: Under-relaxation parameter
"""
struct NonlinearBalance{T<:AbstractFloat} <: AbstractBalanceOperator{T}
    linear_operator::GeostrophicBalance{T}
    max_iterations::Int
    tolerance::T
    relaxation_factor::T
    
    function NonlinearBalance{T}(config::BalanceConfiguration{T};
                                max_iterations::Int = 20,
                                tolerance::T = T(1e-6),
                                relaxation_factor::T = T(0.8)) where T
        linear_op = GeostrophicBalance{T}(config)
        new{T}(linear_op, max_iterations, tolerance, relaxation_factor)
    end
end

"""
    apply_balance_operator(op::GeostrophicBalance{T}, 
                          streamfunction::Matrix{T}) where T

Apply geostrophic balance to compute wind components from streamfunction.

# Mathematical Operation
```
u = -(1/f)∂ψ/∂y, v = (1/f)∂ψ/∂x
```

# Arguments
- `op::GeostrophicBalance{T}`: Geostrophic balance operator
- `streamfunction::Matrix{T}`: Streamfunction field ψ

# Returns
- `Tuple{Matrix{T}, Matrix{T}}`: Geostrophic wind components (u, v)
"""
function apply_balance_operator(op::GeostrophicBalance{T}, 
                               streamfunction::Matrix{T}) where T
    nx, ny = size(streamfunction)
    
    # Flatten for matrix operations
    psi_vec = vec(streamfunction)
    
    # Compute gradients: ∇ψ = (∂ψ/∂x, ∂ψ/∂y)
    grad_x_op, grad_y_op = op.gradient_operators
    dpsi_dx = reshape(grad_x_op * psi_vec, nx, ny)
    dpsi_dy = reshape(grad_y_op * psi_vec, nx, ny)
    
    # Apply geostrophic relations with Coriolis parameter
    f = op.config.coriolis_parameter
    
    # Handle points where f ≈ 0 (near equator)
    f_safe = copy(f)
    equatorial_mask = abs.(f) .< 1e-10
    f_safe[equatorial_mask] .= 1e-10  # Small non-zero value
    
    # Geostrophic wind components
    u_geostrophic = -dpsi_dy ./ f_safe
    v_geostrophic = dpsi_dx ./ f_safe
    
    # Set equatorial winds to zero (no geostrophic approximation)
    u_geostrophic[equatorial_mask] .= zero(T)
    v_geostrophic[equatorial_mask] .= zero(T)
    
    return (u_geostrophic, v_geostrophic)
end

"""
    apply_balance_operator(op::HydrostaticBalance{T}, 
                          temperature::Array{T,3}) where T

Apply hydrostatic balance to compute geopotential from temperature.

# Mathematical Operation  
```
Φ(p) = Φ_surface + ∫[p_surface to p] (RT/p') dp'
```
"""
function apply_balance_operator(op::HydrostaticBalance{T}, 
                               temperature::Array{T,3}) where T
    nx, ny, nz = size(temperature)
    geopotential = zeros(T, nx, ny, nz)
    
    R = op.config.gas_constant
    p_levels = op.pressure_levels
    
    # Integrate hydrostatic equation from surface upward
    for k in 2:nz  # Start from level 2 (level 1 is surface)
        dp = p_levels[k-1] - p_levels[k]  # Pressure thickness
        
        # Average temperature in layer
        temp_avg = (temperature[:, :, k-1] + temperature[:, :, k]) / 2
        
        # Hydrostatic increment: dΦ = -RT/p * dp
        p_avg = (p_levels[k-1] + p_levels[k]) / 2
        geopotential[:, :, k] = geopotential[:, :, k-1] + R * temp_avg * dp / p_avg
    end
    
    return geopotential
end

"""
    apply_balance_operator(op::ThermalWindBalance{T}, 
                          temperature::Array{T,3}) where T

Apply thermal wind balance to compute wind shear from temperature gradients.
"""
function apply_balance_operator(op::ThermalWindBalance{T}, 
                               temperature::Array{T,3}) where T
    nx, ny, nz = size(temperature)
    wind_shear_u = zeros(T, nx, ny, nz-1)
    wind_shear_v = zeros(T, nx, ny, nz-1)
    
    R = op.config.gas_constant
    f = op.config.coriolis_parameter
    dx, dy = op.config.grid_spacing
    
    # Compute temperature gradients and wind shear between levels
    for k in 1:(nz-1)
        # Vertical coordinate difference
        dln_p = log(op.pressure_levels[k]) - log(op.pressure_levels[k+1])
        
        # Horizontal temperature gradients
        dT_dx = zeros(T, nx, ny)
        dT_dy = zeros(T, nx, ny)
        
        # Central differences for interior points
        for j in 2:(ny-1), i in 2:(nx-1)
            dT_dx[i, j] = (temperature[i+1, j, k] - temperature[i-1, j, k]) / (2 * dx)
            dT_dy[i, j] = (temperature[i, j+1, k] - temperature[i, j-1, k]) / (2 * dy)
        end
        
        # Apply boundary conditions (zero gradient)
        dT_dx[1, :] = dT_dx[2, :]
        dT_dx[nx, :] = dT_dx[nx-1, :]
        dT_dy[:, 1] = dT_dy[:, 2]
        dT_dy[:, ny] = dT_dy[:, ny-1]
        
        # Thermal wind relations
        f_safe = copy(f)
        f_safe[abs.(f) .< 1e-10] .= 1e-10
        
        wind_shear_u[:, :, k] = -(R / dln_p) * dT_dy ./ f_safe
        wind_shear_v[:, :, k] = (R / dln_p) * dT_dx ./ f_safe
        
        # Zero shear at equator
        equatorial_mask = abs.(f) .< 1e-10
        wind_shear_u[equatorial_mask, k] .= zero(T)
        wind_shear_v[equatorial_mask, k] .= zero(T)
    end
    
    return (wind_shear_u, wind_shear_v)
end

"""
    apply_balance_operator(op::NonlinearBalance{T}, initial_guess::Dict) where T

Apply nonlinear balance through iterative solution.
"""
function apply_balance_operator(op::NonlinearBalance{T}, initial_guess::Dict) where T
    # Extract fields from initial guess
    vorticity = initial_guess[:vorticity]
    divergence = initial_guess[:divergence]
    
    # Iterative solution of nonlinear balance equation
    streamfunction = copy(initial_guess[:streamfunction])  
    velocity_potential = copy(initial_guess[:velocity_potential])
    
    for iter in 1:op.max_iterations
        # Compute current wind field
        u_geo, v_geo = apply_balance_operator(op.linear_operator, streamfunction)
        
        # Compute residual vorticity
        current_vorticity = compute_vorticity(u_geo, v_geo, op.linear_operator.config)
        residual = vorticity - current_vorticity
        
        # Check convergence
        residual_norm = norm(residual)
        if residual_norm < op.tolerance
            break
        end
        
        # Update streamfunction with relaxation
        correction = solve_poisson_equation(op.linear_operator.laplacian_operator, residual)
        streamfunction .+= op.relaxation_factor * correction
    end
    
    # Compute final balanced wind field
    u_balanced, v_balanced = apply_balance_operator(op.linear_operator, streamfunction)
    
    return Dict(:u => u_balanced, :v => v_balanced, :streamfunction => streamfunction)
end

"""
    compute_balance_residual(op::AbstractBalanceOperator{T}, fields::Dict) where T

Compute residual of balance constraint for diagnostic purposes.
"""
function compute_balance_residual(op::GeostrophicBalance{T}, fields::Dict) where T
    u, v = fields[:u], fields[:v]
    
    # Compute vorticity from wind field
    computed_vorticity = compute_vorticity(u, v, op.config)
    
    # Compute vorticity from streamfunction (if available)
    if haskey(fields, :streamfunction)
        psi = fields[:streamfunction]
        expected_vorticity = reshape(op.laplacian_operator * vec(psi), size(psi))
        residual = computed_vorticity - expected_vorticity
    else
        # Use geostrophic relationship
        f = op.config.coriolis_parameter
        geostrophic_residual = f .* computed_vorticity  # Simplified
        residual = geostrophic_residual
    end
    
    return residual
end

"""
    create_gradient_operators(::Type{T}, nx::Int, ny::Int, 
                            grid_spacing::Tuple{T,T}) where T

Create finite difference gradient operators for 2D grid.
"""
function create_gradient_operators(::Type{T}, nx::Int, ny::Int, 
                                 grid_spacing::Tuple{T,T}) where T
    dx, dy = grid_spacing
    n_points = nx * ny
    
    # X-direction gradient operator (central differences)
    I_x = Int[]
    J_x = Int[]
    V_x = T[]
    
    for j in 1:ny, i in 1:nx
        idx = (j-1) * nx + i
        
        if i == 1  # Forward difference at left boundary
            append!(I_x, [idx, idx])
            append!(J_x, [idx, idx + 1])
            append!(V_x, [-1/(dx), 1/(dx)])
        elseif i == nx  # Backward difference at right boundary
            append!(I_x, [idx, idx])
            append!(J_x, [idx - 1, idx])
            append!(V_x, [-1/(dx), 1/(dx)])
        else  # Central difference in interior
            append!(I_x, [idx, idx, idx])
            append!(J_x, [idx - 1, idx, idx + 1])
            append!(V_x, [-1/(2*dx), 0, 1/(2*dx)])
        end
    end
    
    grad_x = sparse(I_x, J_x, V_x, n_points, n_points)
    
    # Y-direction gradient operator
    I_y = Int[]
    J_y = Int[]
    V_y = T[]
    
    for j in 1:ny, i in 1:nx
        idx = (j-1) * nx + i
        
        if j == 1  # Forward difference at bottom boundary
            append!(I_y, [idx, idx])
            append!(J_y, [idx, idx + nx])
            append!(V_y, [-1/(dy), 1/(dy)])
        elseif j == ny  # Backward difference at top boundary
            append!(I_y, [idx, idx])
            append!(J_y, [idx - nx, idx])
            append!(V_y, [-1/(dy), 1/(dy)])
        else  # Central difference in interior
            append!(I_y, [idx, idx, idx])
            append!(J_y, [idx - nx, idx, idx + nx])
            append!(V_y, [-1/(2*dy), 0, 1/(2*dy)])
        end
    end
    
    grad_y = sparse(I_y, J_y, V_y, n_points, n_points)
    
    return grad_x, grad_y
end

"""
    create_laplacian_operator(::Type{T}, nx::Int, ny::Int, 
                            grid_spacing::Tuple{T,T}) where T

Create 2D Laplacian operator using finite differences.
"""
function create_laplacian_operator(::Type{T}, nx::Int, ny::Int, 
                                 grid_spacing::Tuple{T,T}) where T
    dx, dy = grid_spacing
    n_points = nx * ny
    
    I = Int[]
    J = Int[]
    V = T[]
    
    for j in 1:ny, i in 1:nx
        idx = (j-1) * nx + i
        
        # Interior 5-point stencil: ∇²u = (u_{i-1,j} + u_{i+1,j} - 2u_{i,j})/dx² + 
        #                                 (u_{i,j-1} + u_{i,j+1} - 2u_{i,j})/dy²
        
        coeffs = T[]
        indices = Int[]
        
        # X-direction second derivative
        if i > 1  # Left neighbor
            push!(indices, idx - 1)
            push!(coeffs, 1/(dx^2))
        end
        
        if i < nx  # Right neighbor
            push!(indices, idx + 1)
            push!(coeffs, 1/(dx^2))
        end
        
        # Y-direction second derivative  
        if j > 1  # Bottom neighbor
            push!(indices, idx - nx)
            push!(coeffs, 1/(dy^2))
        end
        
        if j < ny  # Top neighbor
            push!(indices, idx + nx)
            push!(coeffs, 1/(dy^2))
        end
        
        # Central point
        central_coeff = -2/(dx^2) - 2/(dy^2)
        
        # Adjust for boundary conditions (Neumann: zero gradient)
        if i == 1 || i == nx
            central_coeff += 1/(dx^2)  # One less x-neighbor
        end
        if j == 1 || j == ny
            central_coeff += 1/(dy^2)  # One less y-neighbor
        end
        
        push!(indices, idx)
        push!(coeffs, central_coeff)
        
        # Add all entries for this point
        append!(I, fill(idx, length(indices)))
        append!(J, indices)
        append!(V, coeffs)
    end
    
    return sparse(I, J, V, n_points, n_points)
end

"""
    compute_vorticity(u::Matrix{T}, v::Matrix{T}, config::BalanceConfiguration{T}) where T

Compute relative vorticity: ζ = ∂v/∂x - ∂u/∂y.
"""
function compute_vorticity(u::Matrix{T}, v::Matrix{T}, config::BalanceConfiguration{T}) where T
    nx, ny = size(u)
    dx, dy = config.grid_spacing
    
    vorticity = zeros(T, nx, ny)
    
    # Central differences in interior
    for j in 2:(ny-1), i in 2:(nx-1)
        dvdx = (v[i+1, j] - v[i-1, j]) / (2*dx)
        dudy = (u[i, j+1] - u[i, j-1]) / (2*dy)
        vorticity[i, j] = dvdx - dudy
    end
    
    # Boundary conditions (extend interior values)
    vorticity[1, :] = vorticity[2, :]
    vorticity[nx, :] = vorticity[nx-1, :]
    vorticity[:, 1] = vorticity[:, 2]
    vorticity[:, ny] = vorticity[:, ny-1]
    
    return vorticity
end

"""
    Helper functions for additional balance operations
"""

create_divergence_operator(grad_x, grad_y) = grad_x + grad_y  # Simplified

create_vertical_derivative_operator(::Type{T}, pressure_levels::Vector{T}) where T = 
    sparse(I, length(pressure_levels), length(pressure_levels))  # Placeholder

create_thermal_wind_operator(::Type{T}, config, pressure_levels) where T = 
    sparse(I, length(pressure_levels)-1, length(pressure_levels)-1)  # Placeholder

solve_poisson_equation(laplacian_op, rhs) = laplacian_op \ vec(rhs)  # Direct solve

end # module BalanceOperators